import json, time
from openai import OpenAI

start_time = time.time()

client = OpenAI(
    api_key="ollama",
    base_url="http://localhost:11434/v1",
)

system_prompt = """
用户将给你一段病例文本，你需要从中提取用户性别，症状。
最后以JSON的形式输出，格式如下：
{
    "gender": "<性别>返回male或female",
    "symptom": "<症状>返回中文",
    "birth": "<出生日期 YY-MM-DD>"
}
如果用户没有提供相关信息，就用未查询到代替。
"""

user_prompt = """
1.患者老年女性，1965年2月3日生人，因“反复颈痛伴头晕手麻10年，加重1周”入院；
2.主要症状为颈部疼痛伴头晕、手麻；
3.查体：脊柱无畸形、强直、叩击痛，脊柱生理弯曲存在，第5-7颈椎椎体、棘旁两侧肌肉紧张、压痛（+），右侧臂丛神经牵拉试验（+），压头试验（+）。四肢肌力、肌张力对称无异常。神经反射正常存在，未引出病理反射。
4.辅助检查：2015年9月15日我院颈椎MR：颈椎曲度反曲，椎体顺列，颈3-6椎间隙狭窄；颈椎间盘于T2WI上MR信号减低,颈3～6椎体骨质增生边缘变尖；颈3/4、5/6、6/7椎间盘向后突出，硬膜囊前缘受压，脊髓轻度受压；颈4/5椎间盘向右后方突出，硬膜囊前缘受压，右侧神经根受压；脊髓形态信号未见明显异常。
"""

messages = [{"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}]

response = client.chat.completions.create(
    model="deepseek-r1:7b",
    messages=messages,
    response_format={
        'type': 'json_object'
    },
    temperature=0.3,
)

end_time = time.time()

print("实际文本:", user_prompt)
print("deepseek提取json:\n", json.loads(response.choices[0].message.content))
print("耗时:", (end_time - start_time).__round__(2), "秒")
